"""
字母数字识别模块
负责检测和识别图像中的字母和数字
"""

import cv2
import numpy as np
import os
import logging
import re
from typing import List, Tuple, Dict, Optional, Union
try:
    import pytesseract
    from PIL import Image
    TESSERACT_AVAILABLE = True
except ImportError:
    TESSERACT_AVAILABLE = False
    logging.warning("pytesseract未安装，OCR功能将受限")


class OCRDetector:
    """字母数字识别类
    
    提供字母数字检测、识别和文本提取功能
    
    Attributes:
        config: OCR配置参数
        logger: 日志记录器
        tesseract_available: 是否可用Tesseract OCR
    """
    
    def __init__(self, config: Dict):
        """初始化OCR检测器
        
        Args:
            config (dict): OCR配置参数
        """
        self.config = config
        self.logger = logging.getLogger(__name__)
        self.tesseract_available = TESSERACT_AVAILABLE
        
        # 设置Tesseract路径（如果配置中指定）
        if self.tesseract_available and "tesseract_path" in config:
            pytesseract.pytesseract.tesseract_cmd = config["tesseract_path"]
        
        # 字符识别模式
        self.ocr_config = config.get("ocr_config", "--psm 6 --oem 3")
        
        # 字符白名单（限制识别的字符类型）
        self.char_whitelist = config.get("char_whitelist", "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ")
        
        # 最小字符区域大小
        self.min_char_area = config.get("min_char_area", 100)
        
        # 最大字符区域大小
        self.max_char_area = config.get("max_char_area", 10000)
        
        self.logger.info("OCR检测器初始化完成")
    
    def preprocess_image(self, image: np.ndarray) -> np.ndarray:
        """预处理图像以提高OCR识别率
        
        Args:
            image (numpy.ndarray): 输入图像
            
        Returns:
            numpy.ndarray: 预处理后的图像
        """
        # 转换为灰度图像
        if len(image.shape) == 3:
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        else:
            gray = image.copy()
        
        # 去噪
        denoised = cv2.fastNlMeansDenoising(gray)
        
        # 二值化
        _, binary = cv2.threshold(denoised, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
        
        # 形态学操作
        kernel = np.ones((2, 2), np.uint8)
        processed = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel)
        
        return processed
    
    def detect_text_regions(self, image: np.ndarray) -> List[Tuple[int, int, int, int]]:
        """检测图像中的文本区域
        
        Args:
            image (numpy.ndarray): 输入图像
            
        Returns:
            list: 文本区域列表，每个元素为(x, y, w, h)格式的矩形
        """
        # 预处理图像
        processed = self.preprocess_image(image)
        
        # 查找轮廓
        contours, _ = cv2.findContours(processed, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        # 筛选文本区域
        text_regions = []
        for contour in contours:
            x, y, w, h = cv2.boundingRect(contour)
            area = cv2.contourArea(contour)
            
            # 过滤掉过小或过大的区域
            if self.min_char_area <= area <= self.max_char_area:
                # 检查长宽比，字符通常接近正方形或略高
                aspect_ratio = h / w if w > 0 else 0
                if 0.5 <= aspect_ratio <= 3.0:
                    text_regions.append((x, y, w, h))
        
        # 按x坐标排序，从左到右
        text_regions.sort(key=lambda r: r[0])
        
        return text_regions
    
    def recognize_text(self, image: np.ndarray, text_region: Optional[Tuple[int, int, int, int]] = None) -> str:
        """识别图像中的文本
        
        Args:
            image (numpy.ndarray): 输入图像
            text_region (tuple, optional): 文本区域矩形 (x, y, w, h)
            
        Returns:
            str: 识别的文本
        """
        if not self.tesseract_available:
            self.logger.warning("Tesseract OCR不可用，无法识别文本")
            return ""
        
        # 如果指定了文本区域，裁剪图像
        if text_region:
            x, y, w, h = text_region
            roi = image[y:y+h, x:x+w]
        else:
            roi = image
        
        # 预处理图像
        processed = self.preprocess_image(roi)
        
        try:
            # 使用Tesseract识别文本
            custom_config = f'--psm 6 --oem 3 -c tessedit_char_whitelist={self.char_whitelist}'
            text = pytesseract.image_to_string(processed, config=custom_config)
            
            # 清理识别结果
            text = text.strip()
            text = re.sub(r'[^0-9A-Za-z]', '', text)
            
            return text
        except Exception as e:
            self.logger.error(f"文本识别失败: {e}")
            return ""
    
    def recognize_digits(self, image: np.ndarray, text_region: Optional[Tuple[int, int, int, int]] = None) -> str:
        """识别图像中的数字
        
        Args:
            image (numpy.ndarray): 输入图像
            text_region (tuple, optional): 文本区域矩形 (x, y, w, h)
            
        Returns:
            str: 识别的数字
        """
        if not self.tesseract_available:
            self.logger.warning("Tesseract OCR不可用，无法识别数字")
            return ""
        
        # 如果指定了文本区域，裁剪图像
        if text_region:
            x, y, w, h = text_region
            roi = image[y:y+h, x:x+w]
        else:
            roi = image
        
        # 预处理图像
        processed = self.preprocess_image(roi)
        
        try:
            # 使用Tesseract识别数字
            custom_config = f'--psm 7 --oem 3 -c tessedit_char_whitelist=0123456789'
            text = pytesseract.image_to_string(processed, config=custom_config)
            
            # 清理识别结果，只保留数字
            text = text.strip()
            text = re.sub(r'[^0-9]', '', text)
            
            return text
        except Exception as e:
            self.logger.error(f"数字识别失败: {e}")
            return ""
    
    def recognize_letters(self, image: np.ndarray, text_region: Optional[Tuple[int, int, int, int]] = None) -> str:
        """识别图像中的字母
        
        Args:
            image (numpy.ndarray): 输入图像
            text_region (tuple, optional): 文本区域矩形 (x, y, w, h)
            
        Returns:
            str: 识别的字母
        """
        if not self.tesseract_available:
            self.logger.warning("Tesseract OCR不可用，无法识别字母")
            return ""
        
        # 如果指定了文本区域，裁剪图像
        if text_region:
            x, y, w, h = text_region
            roi = image[y:y+h, x:x+w]
        else:
            roi = image
        
        # 预处理图像
        processed = self.preprocess_image(roi)
        
        try:
            # 使用Tesseract识别字母
            custom_config = f'--psm 7 --oem 3 -c tessedit_char_whitelist=ABCDEFGHIJKLMNOPQRSTUVWXYZ'
            text = pytesseract.image_to_string(processed, config=custom_config)
            
            # 清理识别结果，只保留字母
            text = text.strip().upper()
            text = re.sub(r'[^A-Z]', '', text)
            
            return text
        except Exception as e:
            self.logger.error(f"字母识别失败: {e}")
            return ""
    
    def extract_all_text(self, image: np.ndarray) -> List[Dict]:
        """提取图像中的所有文本
        
        Args:
            image (numpy.ndarray): 输入图像
            
        Returns:
            list: 文本信息列表，每个元素包含位置、文本、类型等信息
        """
        # 检测文本区域
        text_regions = self.detect_text_regions(image)
        
        results = []
        for region in text_regions:
            x, y, w, h = region
            
            # 识别文本
            text = self.recognize_text(image, region)
            if not text:
                continue
            
            # 判断文本类型
            if text.isdigit():
                text_type = "digit"
            elif text.isalpha():
                text_type = "letter"
            else:
                text_type = "mixed"
            
            results.append({
                "position": (x + w // 2, y + h // 2),
                "rect": region,
                "text": text,
                "type": text_type,
                "length": len(text)
            })
        
        return results
    
    def draw_text_detections(self, image: np.ndarray, text_results: List[Dict]) -> np.ndarray:
        """在图像上绘制文本检测结果
        
        Args:
            image (numpy.ndarray): 输入图像
            text_results (list): 文本识别结果列表
            
        Returns:
            numpy.ndarray: 绘制了检测结果的图像
        """
        result_image = image.copy()
        
        for result in text_results:
            x, y, w, h = result["rect"]
            text = result["text"]
            text_type = result["type"]
            
            # 根据文本类型选择颜色
            if text_type == "digit":
                color = (0, 0, 255)  # 红色
            elif text_type == "letter":
                color = (0, 255, 0)  # 绿色
            else:
                color = (255, 0, 0)  # 蓝色
            
            # 绘制矩形框
            cv2.rectangle(result_image, (x, y), (x + w, y + h), color, 2)
            
            # 绘制文本
            cv2.putText(result_image, text, (x, y - 10), 
                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2)
            
            # 绘制中心点
            center = result["position"]
            cv2.circle(result_image, center, 3, color, -1)
        
        return result_image
    
    def find_text_pattern(self, image: np.ndarray, pattern: str) -> List[Dict]:
        """查找特定的文本模式
        
        Args:
            image (numpy.ndarray): 输入图像
            pattern (str): 要查找的文本模式
            
        Returns:
            list: 匹配的文本信息列表
        """
        # 提取所有文本
        all_text = self.extract_all_text(image)
        
        # 查找匹配的文本
        matches = []
        for text_info in all_text:
            text = text_info["text"]
            if re.search(pattern, text, re.IGNORECASE):
                matches.append(text_info)
        
        return matches
    
    def read_barcode(self, image: np.ndarray) -> Optional[str]:
        """读取条形码
        
        Args:
            image (numpy.ndarray): 输入图像
            
        Returns:
            str: 条形码内容，如果未检测到返回None
        """
        try:
            import pyzbar.pyzbar as pyzbar
        except ImportError:
            self.logger.warning("pyzbar未安装，无法读取条形码")
            return None
        
        # 转换为灰度图像
        if len(image.shape) == 3:
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        else:
            gray = image
        
        # 检测条形码
        barcodes = pyzbar.decode(gray)
        
        if not barcodes:
            return None
        
        # 返回第一个条形码的内容
        return barcodes[0].data.decode('utf-8')
    
    def read_qrcode(self, image: np.ndarray) -> Optional[str]:
        """读取二维码
        
        Args:
            image (numpy.ndarray): 输入图像
            
        Returns:
            str: 二维码内容，如果未检测到返回None
        """
        try:
            import pyzbar.pyzbar as pyzbar
        except ImportError:
            self.logger.warning("pyzbar未安装，无法读取二维码")
            return None
        
        # 转换为灰度图像
        if len(image.shape) == 3:
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        else:
            gray = image
        
        # 检测二维码
        qrcodes = pyzbar.decode(gray, symbols=[pyzbar.ZBarSymbol.QRCODE])
        
        if not qrcodes:
            return None
        
        # 返回第一个二维码的内容
        return qrcodes[0].data.decode('utf-8')
    
    def verify_text(self, image: np.ndarray, expected_text: str, text_region: Optional[Tuple[int, int, int, int]] = None) -> bool:
        """验证图像中的文本是否与期望文本匹配
        
        Args:
            image (numpy.ndarray): 输入图像
            expected_text (str): 期望的文本
            text_region (tuple, optional): 文本区域矩形 (x, y, w, h)
            
        Returns:
            bool: 匹配返回True，不匹配返回False
        """
        # 识别文本
        recognized_text = self.recognize_text(image, text_region)
        
        # 比较文本
        return recognized_text.upper() == expected_text.upper()